# file: download_video.py
import os
import time
import datetime
import shutil
import subprocess
from ctypes import *
import glob

# ============================================================================
# 1. 配置参数 (您修改的路径已保留)
# ============================================================================
# NVR 连接信息
NVR_IP = b"192.168.1.30"
NVR_PORT = 80
NVR_USERNAME = b"admin"
NVR_PASSWORD = b"Nanjing@8625"

# 摄像头信息列表
CAMERAS = [
    {"code": "210235C8Q33254000011", "channel_id": 1},
    {"code": "210235C8Q33255000005", "channel_id": 2},
    {"code": "210235C8Q33255000006", "channel_id": 3},
]

# 下载与处理参数
DOWNLOAD_DURATION_SECONDS = 60
LOCAL_VIDEO_SAVE_DIR = "./downloaded_videos_temp"

# 您修改后的图片路径
FINAL_IMAGE_ROOT = "/home/smf/image"
FINAL_TS_ROOT = "/mnt/nfs_share/"

# SDK库文件路径
SDK_LIB_PATH = "/root/egs_nvr-master_x86/lib"

# ============================================================================
# 2. SDK 初始化和定义 (DevClass.download方法已更新)
# ============================================================================
net_sdk_dll_path = os.path.join(SDK_LIB_PATH, "libNetDEVSDK.so")
os.environ['LD_LIBRARY_PATH'] = f"{SDK_LIB_PATH}:{os.environ.get('LD_LIBRARY_PATH', '')}"

try:
    from netdevsdk import *
except ImportError:
    print("警告：无法从'netdevsdk'导入类型定义，将使用占位符。")
    class c_void_p: pass
    class POINTER: pass
    class c_int: pass
    class c_char_p: pass
    class c_longlong: pass
    class c_int64: pass
    class c_int32: pass
    class byref: pass
    class pointer: pass
    def load_library(path): return cdll.LoadLibrary(path)
    class NETDEV_DEVICE_LOGIN_INFO_S: pass
    class NETDEV_SELOG_INFO_S: pass
    class NETDEV_PLAYBACKCOND_S: pass
    class NETDEV_VIDEO_CHL_DETAIL_INFO_S: pass


class DevClass:
    def __init__(self):
        self.egsSdk = self.load_sdk()
        if self.egsSdk:
            self.setup_sdk_functions()
        self.iUserID = -1
    def load_sdk(self):
        egs_sdk = None
        try:
            egs_sdk = load_library(net_sdk_dll_path)
            print(f"动态库[{net_sdk_dll_path}]加载成功")
        except Exception as e:
            print(f"动态库[{net_sdk_dll_path}]加载失败:{e}")
        return egs_sdk
    def setup_sdk_functions(self):
        self.egsSdk.NETDEV_Login_V30.restype = c_longlong
        self.egsSdk.NETDEV_GetFileByTime.argtypes = [c_longlong, POINTER(NETDEV_PLAYBACKCOND_S), c_char_p, c_int]
        self.egsSdk.NETDEV_GetFileByTime.restype = c_longlong
        self.egsSdk.NETDEV_Logout.argtypes = [c_longlong]
        self.egsSdk.NETDEV_PlayBackControl.argtypes = [c_longlong, c_int, POINTER(c_int64)]
        self.egsSdk.NETDEV_StopGetFile.argtypes = [c_longlong]
    def sdk_init(self):
        res = self.egsSdk.NETDEV_Init()
        if res != 0: print(f"初始化sdk失败，错误码: {res}")
        else: print("初始化sdk成功")
    def login(self, ip, port, username, password):
        login_info = NETDEV_DEVICE_LOGIN_INFO_S()
        login_info.szIPAddr = ip
        login_info.dwPort = port
        login_info.szUserName = username
        login_info.szPassword = password
        login_info.dwLoginProto = 0
        login_info.dwDeviceType = 0
        se_log_info = NETDEV_SELOG_INFO_S()
        self.iUserID = self.egsSdk.NETDEV_Login_V30(pointer(login_info), pointer(se_log_info))
        if self.iUserID <= 0:
            print("登陆失败，错误码：", self.egsSdk.NETDEV_GetLastError())
            return False
        print(f"登录成功，用户句柄 iUserID: {hex(self.iUserID)}")
        return True

    # MODIFIED: 更新了此方法以提高下载的稳定性
    def download(self, channel_id, begin_time, end_time, save_filename):
        save_dir = os.path.dirname(save_filename)
        if not os.path.exists(save_dir): os.makedirs(save_dir)

        playback_cond = NETDEV_PLAYBACKCOND_S()
        playback_cond.dwChannelID = channel_id
        playback_cond.tBeginTime = begin_time
        playback_cond.tEndTime = end_time
        playback_cond.dwLinkMode = 1
        playback_cond.dwDownloadSpeed = 8  # 下载倍速
        print(f"开始下载，用户句柄 iUserID: {hex(self.iUserID)}")
        get_file = self.egsSdk.NETDEV_GetFileByTime(self.iUserID, byref(playback_cond), save_filename.encode('utf-8'), 0)

        if get_file <= 0:
            print("启动按时间下载失败，错误码:", self.egsSdk.NETDEV_GetLastError())
            return False
            
        print(f"获取下载句柄成功: {hex(get_file)}，开始查询进度...")
        download_completed = False
        
        # 增加超时机制
        timeout_start = time.time()
        max_wait_time = DOWNLOAD_DURATION_SECONDS + 30  # 视频时长 + 30秒缓冲

        try:
            play_time = c_int64(0)
            while True:
                # 1. 检查是否超时
                if time.time() - timeout_start > max_wait_time:
                    print("下载超时，任务中断。")
                    break

                # 2. 获取下载进度，并检查函数返回值
                res = self.egsSdk.NETDEV_PlayBackControl(get_file, 3, byref(play_time))
                
                # 3. 检查SDK是否报告流结束或错误
                if res == -1:
                    print("SDK报告下载流已结束或发生错误。")
                    # 检查是否是因下载完成而结束
                    if play_time.value >= end_time:
                        download_completed = True
                    break

                # 4. 检查是否下载完成
                if play_time.value >= end_time:
                    print("当前下载到：100.00%")
                    download_completed = True
                    break
                
                # 5. 仅当进度有效时才打印，避免初始时打印无意义的进度
                # if play_time.value > 0 and end_time > begin_time:
                #     progress = ((play_time.value - begin_time) / (end_time - begin_time)) * 100
                #     print(f"当前下载到：{progress:.2f}%%")
                
                time.sleep(1) # 每次查询间隔1秒
        finally:
            print(f"停止下载句柄: {hex(get_file)}")
            self.egsSdk.NETDEV_StopGetFile(get_file)
            print("下载任务结束。")
        
        return download_completed
        
    def logout(self):
        if self.iUserID <= 0: return
        res = self.egsSdk.NETDEV_Logout(self.iUserID)
        if res != 0: print("登出失败，错误码：", self.egsSdk.NETDEV_GetLastError())
        else: print("登出成功")
    def sdk_cleanup(self):
        res = self.egsSdk.NETDEV_Cleanup()
        if res != 0: print(f"释放sdk失败，错误码: {res}")
        else: print("释放sdk成功")
    def __del__(self):
        self.logout()
        self.sdk_cleanup()

# ============================================================================
# 3. 视频处理函数 (已按新要求修改)
# ============================================================================

# MODIFIED: 函数1 - 将MP4抽帧为PNG图片
# 现在会返回第一个PNG的路径和其所在目录的路径，用于后续生成缩略图
def extract_images(source_mp4_path, camera_code, start_timestamp):
    """将MP4视频抽帧为PNG图片，并返回首张图片的路径和图片目录。"""
    print("\n--- 开始处理 [抽帧为PNG] ---")
    if not os.path.exists(source_mp4_path):
        print(f"[图片处理失败] 源文件不存在: {source_mp4_path}")
        return None, None # 返回None表示失败
    try:
        start_dt = datetime.datetime.fromtimestamp(start_timestamp)
        date_path = start_dt.strftime('%Y/%m/%d')
        hour_path = start_dt.strftime('%H')
        image_folder_name = start_dt.strftime('%H%M%S')
        final_dir = os.path.join(FINAL_IMAGE_ROOT, camera_code, date_path, hour_path, image_folder_name)
        os.makedirs(final_dir, exist_ok=True)
        print(f"图片将保存到: {final_dir}")

        output_pattern = os.path.join(final_dir, "%06d.png")
        first_png_path = os.path.join(final_dir, "000001.png") # 预先定义第一张图的路径
        
        ffmpeg_command = ["ffmpeg", "-i", source_mp4_path, output_pattern]

        print(f"执行ffmpeg图片抽帧命令: {' '.join(ffmpeg_command)}")
        subprocess.run(ffmpeg_command, check=True, capture_output=True, text=True)
        print(f"[图片处理成功] 视频已抽帧为图片。")
        
        # 返回首图路径和目录路径
        return first_png_path, final_dir
    except Exception as e:
        print(f"[错误] 图片处理过程中发生异常: {e}")
        return None, None

# MODIFIED: 函数2 - 将MP4分割为TS文件
# 现在会返回生成的TS文件路径列表，用于后续生成M3U8文件
def convert_mp4_to_ts(source_mp4_path, camera_code, start_timestamp):
    """将MP4视频分割为TS文件，并返回TS文件路径列表。"""
    print("\n--- 开始处理 [分割为TS] ---")
    if not os.path.exists(source_mp4_path):
        print(f"[TS处理失败] 源文件不存在: {source_mp4_path}")
        return [] # 返回空列表表示失败

    start_dt = datetime.datetime.fromtimestamp(start_timestamp)
    date_path = start_dt.strftime('%Y/%m/%d')
    hour_path = start_dt.strftime('%H')
    final_dir = os.path.join(FINAL_TS_ROOT, camera_code, date_path, hour_path)
    
    final_ts_paths = [] # 用于存储最终的TS文件路径
    try:
        os.makedirs(final_dir, exist_ok=True)
        print(f"TS文件将保存到: {final_dir}")

        temp_output_pattern = os.path.join(final_dir, f"temp_{start_timestamp}_%02d.ts")
        
        ffmpeg_command = [
            "ffmpeg", "-i", source_mp4_path, "-c", "copy", "-map", "0",
            "-f", "segment", "-segment_time", "10", "-segment_format", "mpegts",
            "-reset_timestamps", "1", temp_output_pattern
        ]

        print(f"执行ffmpeg TS分割命令 (生成临时文件): {' '.join(ffmpeg_command)}")
        subprocess.run(ffmpeg_command, check=True, capture_output=True, text=True)
        
        print("开始重命名TS文件...")
        temp_files = sorted(glob.glob(os.path.join(final_dir, f"temp_{start_timestamp}_*.ts")))
        
        segment_duration = 10
        for i, temp_file_path in enumerate(temp_files):
            current_segment_timestamp = start_timestamp + (i * segment_duration)
            new_filename = datetime.datetime.fromtimestamp(current_segment_timestamp).strftime('%H%M%S') + ".ts"
            new_file_path = os.path.join(final_dir, new_filename)
            os.rename(temp_file_path, new_file_path)
            final_ts_paths.append(new_file_path) # 将新路径添加到列表中
            print(f"  '{os.path.basename(temp_file_path)}' -> '{new_filename}'")

        print(f"[TS处理成功] 视频已分割并重命名为TS文件。")
        return final_ts_paths
    except Exception as e:
        temp_files_to_clean = glob.glob(os.path.join(final_dir, f"temp_{start_timestamp}_*.ts"))
        for f in temp_files_to_clean: os.remove(f)
        print(f"[错误] TS文件处理过程中发生异常: {e}")
        return []

# NEW: 函数3 - 从第一张PNG创建JPG缩略图
def create_jpg_thumbnail(first_png_path, png_folder_path, camera_code, start_timestamp):
    """从PNG创建JPG缩略图并保存到指定的source目录。"""
    print("\n--- 开始处理 [生成JPG缩略图] ---")
    if not first_png_path or not os.path.exists(first_png_path):
        print(f"[缩略图处理失败] 源PNG文件不存在: {first_png_path}")
        return False
        
    try:
        # 1. 构建目标目录
        start_dt = datetime.datetime.fromtimestamp(start_timestamp)
        dest_dir = os.path.join(
            FINAL_TS_ROOT, # 复用TS根目录 /mnt/nfs_share/
            camera_code,
            start_dt.strftime('%Y/%m/%d'),
            'source'
        )
        os.makedirs(dest_dir, exist_ok=True)

        # 2. 构建目标文件名（使用PNG所在文件夹的名字）
        jpg_filename = os.path.basename(png_folder_path) + "_60.jpg"
        dest_jpg_path = os.path.join(dest_dir, jpg_filename)

        # 3. 使用ffmpeg进行转换
        # -q:v 2 是一个很好的JPG质量参数 (范围2-31，越低越好)
        ffmpeg_command = ["ffmpeg", "-y", "-i", first_png_path, "-q:v", "2", dest_jpg_path]
        
        print(f"执行ffmpeg JPG转换命令: {' '.join(ffmpeg_command)}")
        subprocess.run(ffmpeg_command, check=True, capture_output=True, text=True)
        print(f"[缩略图处理成功] JPG已保存到: {dest_jpg_path}")
        return True
    except Exception as e:
        print(f"[错误] JPG缩略图处理过程中发生异常: {e}")
        return False

# NEW: 函数4 - 为TS文件列表创建M3U8播放列表
def create_m3u8_playlist(ts_file_list, camera_code, start_timestamp):
    """为TS文件列表创建M3U8播放列表。"""
    print("\n--- 开始处理 [生成M3U8文件] ---")
    if not ts_file_list:
        print("[M3U8处理失败] TS文件列表为空。")
        return False
        
    try:
        # 1. 构建目标目录 (与JPG缩略图目录相同)
        start_dt = datetime.datetime.fromtimestamp(start_timestamp)
        dest_dir = os.path.join(
            FINAL_TS_ROOT,
            camera_code,
            start_dt.strftime('%Y/%m/%d'),
            'source'
        )
        os.makedirs(dest_dir, exist_ok=True)
        
        # 2. 构建目标文件名
        m3u8_filename = start_dt.strftime('%H%M%S') + "_60.m3u8"
        m3u8_filepath = os.path.join(dest_dir, m3u8_filename)

        # 3. 生成M3U8内容
        m3u8_content = [
            "#EXTM3U",
            "#EXT-X-VERSION:3",
            "#EXT-X-TARGETDURATION:10", # TS分片时长
            "#EXT-X-MEDIA-SEQUENCE:0"
        ]
        
        for ts_path in ts_file_list:
            m3u8_content.append("#EXTINF:10.000,")
            # 关键：M3U8中的路径应该是相对于M3U8文件的相对路径
            # 因为TS和M3U8不在同一目录，需要使用 `../小时/文件名.ts`
            ts_filename = os.path.basename(ts_path)
            hour_folder = os.path.basename(os.path.dirname(ts_path))
            relative_path = f"../{hour_folder}/{ts_filename}"
            m3u8_content.append(relative_path)
            
        m3u8_content.append("#EXT-X-ENDLIST")

        # 4. 写入文件
        with open(m3u8_filepath, 'w') as f:
            f.write('\n'.join(m3u8_content))
            
        print(f"[M3U8处理成功] 播放列表已保存到: {m3u8_filepath}")
        return True
    except Exception as e:
        print(f"[错误] M3U8文件处理过程中发生异常: {e}")
        return False

# ============================================================================
# 4. 主循环 (已修改为测试模式：下载3个视频并保留)
# ============================================================================
if __name__ == '__main__':
    dev = DevClass()
    dev.sdk_init()

    if not dev.login(NVR_IP, NVR_PORT, NVR_USERNAME, NVR_PASSWORD):
        exit()

    # 起始时间可以根据您的需要修改
    start_dt = datetime.datetime(2025, 7, 22, 15, 20, 0)
    start_time = int(start_dt.timestamp())

    print(f"测试任务启动，起始时间设置为: {start_dt.strftime('%Y-%m-%d %H:%M:%S')}")
    print("将为每个摄像头下载3个一分钟的视频片段，并保留MP4文件。")

    # 使用 for 循环代替 while True，精确执行3次
    for i in range(3):
        try:
            # 计算当前时间片的结束时间
            end_time = start_time + DOWNLOAD_DURATION_SECONDS
            
            print(f"\n--- [第 {i+1}/3 次] 开始处理时间片: {start_time} ({datetime.datetime.fromtimestamp(start_time)}) ---")
            
            for camera in CAMERAS:
                camera_code = camera["code"]
                channel_id = camera["channel_id"]
                print(f"\n>>>>>> 处理摄像头: {camera_code} (通道: {channel_id}) <<<<<<")
                
                nvr_ip_str = NVR_IP.decode('utf-8')
                temp_dir = os.path.join(LOCAL_VIDEO_SAVE_DIR, nvr_ip_str)
                temp_video_filename = os.path.join(temp_dir, f"{camera_code}_{start_time}.mp4")
                
                if dev.download(channel_id, start_time, end_time, temp_video_filename):
                    print(f"\n视频片段 {temp_video_filename} 下载成功，开始后续处理。")
                    
                    # 即使在测试中，后续处理流程也会照常运行
                    first_png, png_dir = extract_images(temp_video_filename, camera_code, start_time)
                    ts_file_list = convert_mp4_to_ts(temp_video_filename, camera_code, start_time)

                    if first_png and png_dir:
                        create_jpg_thumbnail(first_png, png_dir, camera_code, start_time)
                    
                    if ts_file_list:
                        create_m3u8_playlist(ts_file_list, camera_code, start_time)

                    # 关键修改：注释掉下面的代码块以保留MP4文件
                    # if os.path.exists(temp_video_filename):
                    #     os.remove(temp_video_filename)
                    #     print(f"\n已清理临时文件: {temp_video_filename}")
                    print(f"\n[测试模式] MP4文件已保留: {temp_video_filename}")

                else:
                    print(f"摄像头 {camera_code} 下载失败或未下载完整，将跳过此时间片。")
            
            # 更新下一次下载的开始时间，确保下载连续的三个时间片
            start_time = end_time

        except KeyboardInterrupt:
            print("\n接收到手动停止命令，正在清理和退出...")
            break
        except Exception as e:
            print(f"主循环发生未知错误: {e}。")
            break
            
    print("\n\n测试任务完成。")
    print("程序已退出。")